import uuid
from datetime import datetime
from collections import deque


class HistoryMemory:
    """
    管理对话历史记录的类，支持用户隔离、会话ID和上下文检索。
    """

    def __init__(self, max_sessions=10, max_turns_per_session=20):
        """
        初始化 HistoryMemory 实例。

        Args:
            max_sessions (int): 每个用户保留的最大会话数。默认为 10。
            max_turns_per_session (int): 每个会话保留的最大对话轮数。默认为 20。
        """
        self.max_sessions = max_sessions
        self.max_turns_per_session = max_turns_per_session

        # 存储所有用户的会话 {user_id: {session_id: deque of turns}}
        self.user_sessions = {}
        # 当前活动会话信息
        self.current_user_id = None
        self.current_session_id = None

    def _get_or_create_session(self, session_id=None, user_id=None):
        """获取现有会话或创建新会话（支持用户隔离）"""
        # 使用当前用户ID如果未指定
        user_id = user_id or self.current_user_id
        if not user_id:
            raise ValueError("用户ID不能为空")

        # 确保用户会话容器存在
        if user_id not in self.user_sessions:
            self.user_sessions[user_id] = {}

        # 生成新会话ID如果未指定
        if session_id is None:
            session_id = str(uuid.uuid4())

        # 确保会话存在
        if session_id not in self.user_sessions[user_id]:
            self.user_sessions[user_id][session_id] = deque(maxlen=self.max_turns_per_session)

            # 管理用户的会话数量
            if len(self.user_sessions[user_id]) > self.max_sessions:
                # 移除最早的会话
                first_key = next(iter(self.user_sessions[user_id]))
                del self.user_sessions[user_id][first_key]

        # 更新当前会话信息
        self.current_user_id = user_id
        self.current_session_id = session_id
        return self.user_sessions[user_id][session_id]

    def start_new_session(self, user_id=None):
        """
        开始一个新会话

        Args:
            user_id (str, optional): 用户ID，用于隔离不同用户的会话

        Returns:
            str: 新会话的ID
        """
        new_session_id = str(uuid.uuid4())
        self._get_or_create_session(new_session_id, user_id)
        return new_session_id

    def add_turn(self, user_input: str, agent_response: str, session_id=None, user_id=None):
        """
        添加一轮对话到指定或当前会话的历史记录中。

        Args:
            user_input (str): 用户的输入。
            agent_response (str): Agent 的回复。
            session_id (str, optional): 会话 ID。如果为 None，则使用当前会话。
            user_id (str, optional): 用户ID，用于隔离不同用户的会话
        """
        session_history = self._get_or_create_session(session_id, user_id)
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        turn = {
            "id": str(uuid.uuid4()),  # 为每轮对话分配唯一 ID
            "timestamp": timestamp,
            "user": user_input,
            "agent": agent_response
        }
        session_history.append(turn)

    def get_history(self, session_id=None, user_id=None) -> list:
        """
        获取指定或当前会话的对话历史记录列表。

        Args:
            session_id (str, optional): 会话 ID。如果为 None，则使用当前会话。
            user_id (str, optional): 用户ID，用于隔离不同用户的会话

        Returns:
            list: 包含对话历史的字典列表。
        """
        user_id = user_id or self.current_user_id
        session_id = session_id or self.current_session_id

        if user_id and session_id and user_id in self.user_sessions and session_id in self.user_sessions[user_id]:
            return list(self.user_sessions[user_id][session_id])
        return []

    def get_context(self, session_id=None, user_id=None, last_n=5) -> str:
        """
        获取指定或当前会话的上下文字符串，用于构建 Prompt。

        Args:
            session_id (str, optional): 会话 ID。
            user_id (str, optional): 用户ID，用于隔离不同用户的会话
            last_n (int): 包含最近的 N 轮对话。

        Returns:
            str: 格式化的上下文字符串。
        """
        history = self.get_history(session_id, user_id)
        recent_history = history[-last_n:] if len(history) >= last_n else history

        context_parts = []
        for turn in recent_history:
            context_parts.append(f"用户: {turn['user']}")
            context_parts.append(f"助手: {turn['agent']}")

        return "\n".join(context_parts)

    def get_sessions_list(self, user_id=None) -> list:
        """
        获取指定用户的所有会话简要信息列表，按时间倒序排列。

        Args:
            user_id (str, optional): 用户ID，用于隔离不同用户的会话

        Returns:
            list: 包含会话信息的字典列表。
        """
        user_id = user_id or self.current_user_id
        if not user_id or user_id not in self.user_sessions:
            return []

        session_list = []
        for sid, history in self.user_sessions[user_id].items():
            if history:
                # 使用第一轮对话的用户输入作为标题
                first_turn = history[0]
                title = first_turn['user'][:30] + ('...' if len(first_turn['user']) > 30 else '')
                timestamp = first_turn['timestamp']
            else:
                title = "新会话"
                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

            session_list.append({
                "id": sid,
                "title": title,
                "timestamp": timestamp
            })

        # 按时间戳倒序排序
        session_list.sort(key=lambda x: x['timestamp'], reverse=True)
        return session_list

    def clear_session(self, session_id=None, user_id=None):
        """
        清空指定或当前会话的历史记录。

        Args:
            session_id (str, optional): 会话 ID。
            user_id (str, optional): 用户ID
        """
        user_id = user_id or self.current_user_id
        session_id = session_id or self.current_session_id

        if user_id and session_id and user_id in self.user_sessions and session_id in self.user_sessions[user_id]:
            self.user_sessions[user_id][session_id].clear()

    def delete_session(self, session_id, user_id=None):
        """
        删除指定的会话。

        Args:
            session_id (str): 要删除的会话 ID。
            user_id (str, optional): 用户ID
        """
        user_id = user_id or self.current_user_id
        if user_id and user_id in self.user_sessions and session_id in self.user_sessions[user_id]:
            del self.user_sessions[user_id][session_id]
            if self.current_session_id == session_id and self.current_user_id == user_id:
                self.current_session_id = None

    def clear_all_sessions(self, user_id=None):
        """
        清空指定用户的所有会话

        Args:
            user_id (str, optional): 用户ID，如果为None则清空当前用户的会话
        """
        user_id = user_id or self.current_user_id
        if user_id and user_id in self.user_sessions:
            self.user_sessions[user_id].clear()
            if self.current_user_id == user_id:
                self.current_session_id = None

    def set_current_session(self, session_id, user_id=None):
        """
        设置当前活动的会话。

        Args:
            session_id (str): 要设置为当前的会话 ID。
            user_id (str, optional): 用户ID
        """
        user_id = user_id or self.current_user_id
        if user_id and user_id in self.user_sessions and session_id in self.user_sessions[user_id]:
            self.current_user_id = user_id
            self.current_session_id = session_id


# --- 示例用法 ---
if __name__ == "__main__":
    memory = HistoryMemory(max_sessions=3, max_turns_per_session=5)

    # 用户1的会话
    user1_id = "user_123"
    session1_id = memory.start_new_session(user1_id)
    memory.add_turn("你好！", "你好！有什么可以帮助你的吗？", session1_id, user1_id)
    memory.add_turn("今天天气如何？", "今天天气晴朗。", session1_id, user1_id)

    # 用户2的会话
    user2_id = "user_456"
    session2_id = memory.start_new_session(user2_id)
    memory.add_turn("讲个笑话吧。", "为什么电脑去医院？因为它需要‘重启’一下！", session2_id, user2_id)

    # 继续用户1的会话
    memory.add_turn("谢谢！", "不客气！", session1_id, user1_id)

    print(f"--- 用户 {user1_id} 的会话列表 ---")
    print(memory.get_sessions_list(user1_id))

    print(f"\n--- 用户 {user1_id} 的会话 {session1_id} 历史 ---")
    print(memory.get_history(session1_id, user1_id))

    print(f"\n--- 用户 {user2_id} 的会话列表 ---")
    print(memory.get_sessions_list(user2_id))
